from collections import deque
from typing import Any, SupportsFloat
import gymnasium as gym
import numpy as np

class RecordSuccessRateWrapper(gym.Wrapper):
    def __init__(self, env: gym.Env):
        super().__init__(env)
        self.recorder = deque(maxlen=20)
        self.difficulty_level = 0 #由于浮点数的计算问题 -0.1+0.1!=0，我又不想改环境，在这里新加一个变量来记录难度
        self.boundary_0=np.zeros((2,3))#用于限制机器人的运动范围
        self.boundary_0[0]=self.env.unwrapped.target_pos[:3]-0.01
        self.boundary_0[0][2]=self.env.unwrapped.target_pos[2]-0.003
        self.boundary_0[1]=self.env.unwrapped.target_pos[:3]+0.01
        self.boundary_0[1][2]=self.env.unwrapped.target_pos[2]+0.015
        self.boundary_1 = self.env.unwrapped.boundary.copy()
        if self.env.unwrapped.curriculum_difficulty == 0:
            self.env.unwrapped.boundary=self.boundary_0.copy()

    def step(self, action: Any):
        observation, reward, terminated ,truncated ,info = self.env.step(action)
        if terminated:
            self.recorder.append(1)
            info = self.update_info(info)
        elif truncated:
            self.recorder.append(0)
            info = self.update_info(info)         
        return observation, reward, terminated, truncated, info
    
    def update_info(self, info):
        if len(self.recorder)==self.recorder.maxlen:
            info['success_rate'] = sum(self.recorder) / len(self.recorder)
            self.env.unwrapped.success_rate = info['success_rate']
            if self.env.unwrapped.curriculum :
                if self.difficulty_level < 10 and self.env.unwrapped.success_rate>=0.8:
                    difficulty = self.difficulty_level
                    self.difficulty_level += 1
                    self.env.unwrapped.curriculum_difficulty = self.difficulty_level/10
                    for _ in range(15):
                        self.recorder.popleft()
                    if difficulty == 6 and self.difficulty_level == 7:
                        info['switch_demo'] = 1
                    if difficulty == 0 and self.difficulty_level == 1:
                        info['switch_demo'] = 1
                        self.env.unwrapped.boundary=self.boundary_1.copy()
                elif self.difficulty_level > 0 and self.env.unwrapped.success_rate<=0.5:
                    difficulty = self.difficulty_level
                    self.difficulty_level -= 1
                    self.env.unwrapped.curriculum_difficulty = self.difficulty_level / 10
                    for _ in range(15):
                        self.recorder.popleft()
                    if difficulty == 7 and self.difficulty_level == 6:
                        info['switch_demo'] = -1
                    if difficulty == 1 and self.difficulty_level == 0:
                        info['switch_demo'] = -1

        return info

    